package Rijndael;

import org.apache.log4j.Logger;

/**
 * Created by IntelliJ IDEA.
 * User: dima
 * Date: 18.09.2010
 * Time: 19:13:47
 * реализация шифра AES
 */
public class Encoder {
    private static final Logger log = Logger.getLogger(Encoder.class);
    public static final int SIZE = 4;
    private static final int[][] sBox = new int [][]
         {{0x63 ,0x7c ,0x77 ,0x7b ,0xf2 ,0x6b ,0x6f ,0xc5 ,0x30 ,0x01 ,0x67 ,0x2b ,0xfe ,0xd7 ,0xab ,0x76},
          {0xca ,0x82 ,0xc9 ,0x7d ,0xfa ,0x59 ,0x47 ,0xf0 ,0xad ,0xd4 ,0xa2 ,0xaf ,0x9c ,0xa4 ,0x72 ,0xc0},
          {0xb7 ,0xfd ,0x93 ,0x26 ,0x36 ,0x3f ,0xf7 ,0xcc ,0x34 ,0xa5 ,0xe5 ,0xf1 ,0x71 ,0xd8 ,0x31 ,0x15},
          {0x04 ,0xc7 ,0x23 ,0xc3 ,0x18 ,0x96 ,0x05 ,0x9a ,0x07 ,0x12 ,0x80 ,0xe2 ,0xeb ,0x27 ,0xb2 ,0x75},
          {0x09 ,0x83 ,0x2c ,0x1a ,0x1b ,0x6e ,0x5a ,0xa0 ,0x52 ,0x3b ,0xd6 ,0xb3 ,0x29 ,0xe3 ,0x2f ,0x84},
          {0x53 ,0xd1 ,0x00 ,0xed ,0x20 ,0xfc ,0xb1 ,0x5b ,0x6a ,0xcb ,0xbe ,0x39 ,0x4a ,0x4c ,0x58 ,0xcf},
          {0xd0 ,0xef ,0xaa ,0xfb ,0x43 ,0x4d ,0x33 ,0x85 ,0x45 ,0xf9 ,0x02 ,0x7f ,0x50 ,0x3c ,0x9f ,0xa8},
          {0x51 ,0xa3 ,0x40 ,0x8f ,0x92 ,0x9d ,0x38 ,0xf5 ,0xbc ,0xb6 ,0xda ,0x21 ,0x10 ,0xff ,0xf3 ,0xd2},
          {0xcd ,0x0c ,0x13 ,0xec ,0x5f ,0x97 ,0x44 ,0x17 ,0xc4 ,0xa7 ,0x7e ,0x3d ,0x64 ,0x5d ,0x19 ,0x73},
          {0x60 ,0x81 ,0x4f ,0xdc ,0x22 ,0x2a ,0x90 ,0x88 ,0x46 ,0xee ,0xb8 ,0x14 ,0xde ,0x5e ,0x0b ,0xdb},
          {0xe0 ,0x32 ,0x3a ,0x0a ,0x49 ,0x06 ,0x24 ,0x5c ,0xc2 ,0xd3 ,0xac ,0x62 ,0x91 ,0x95 ,0xe4 ,0x79},
          {0xe7 ,0xc8 ,0x37 ,0x6d ,0x8d ,0xd5 ,0x4e ,0xa9 ,0x6c ,0x56 ,0xf4 ,0xea ,0x65 ,0x7a ,0xae ,0x08},
          {0xba ,0x78 ,0x25 ,0x2e ,0x1c ,0xa6 ,0xb4 ,0xc6 ,0xe8 ,0xdd ,0x74 ,0x1f ,0x4b ,0xbd ,0x8b ,0x8a},
          {0x70 ,0x3e ,0xb5 ,0x66 ,0x48 ,0x03 ,0xf6 ,0x0e ,0x61 ,0x35 ,0x57 ,0xb9 ,0x86 ,0xc1 ,0x1d ,0x9e},
          {0xe1 ,0xf8 ,0x98 ,0x11 ,0x69 ,0xd9 ,0x8e ,0x94 ,0x9b ,0x1e ,0x87 ,0xe9 ,0xce ,0x55 ,0x28 ,0xdf},
          {0x8c ,0xa1 ,0x89 ,0x0d ,0xbf ,0xe6 ,0x42 ,0x68 ,0x41 ,0x99 ,0x2d ,0x0f ,0xb0 ,0x54 ,0xbb ,0x16}};

    /**
     * адаптер для функции кодирования, позволяющий принимать на вход строки
     *
     * @param inStr открытый текст (32 символа)
     * @param keyStr ключ длиной 32 символа
     * @return строка шифротекста длиной 32 символа
     */
    public static String encode(String inStr, String keyStr) {
        int[][] in = new int[SIZE][SIZE];
        int[][] key = new int[SIZE][SIZE];
        int strPos = 0;
//        заполнение происходит по строкам
        for (int i = 0; i < SIZE; i++) {
            for (int j = 0; j < SIZE; j++) {
//                преобразуем два подряд идущих символа в байт
//                для открытого текста
                String twoChars = inStr.substring(strPos,strPos+2);
                in[j][i] = Integer.parseInt(twoChars, 16);
//                для ключа
                twoChars = keyStr.substring(strPos,strPos+2);
                key[i][j] = Integer.parseInt(twoChars, 16);
                strPos += 2;
            }
        }
        int [][]cipher = encode(in, key);
//        обратное преобразование таблицы чисел в строку
        return blockToString(cipher,false);
    }

    /**
     * преобразование блока даных в HEX-строчку
     * @param block блок данных
     * @param enableLog если true, результат пойдет в лог
     * @return ___
     */
    protected static String blockToString(int[][] block, boolean enableLog) {
        String blockStr = "";
//        log.debug("blockToString:");
        for (int i = 0; i < SIZE; i++) {
            for (int j = 0; j < SIZE; j++) {
//                преобразуем байт в два подряд идущих символа
                blockStr += Integer.toHexString(block[j][i] >> 4) + Integer.toHexString(block[j][i] & 0xf);
//                blockStr += Integer.toHexString(block[i][j]);
            }
            if(enableLog){
                log.debug(blockStr.substring(blockStr.length()-8,blockStr.length()));
            }
        }
        return blockStr;
    }

    /**
     * функция кодирования AES
     *
     * @param in Открытый текст в виде таблицы чисел 4х4
     * @param key ключ в том же виде что и открытый текст
     * @return шифротекст 4х4
     */
    public static int[][] encode(int[][] in, int[][] key) {
        int state [][] = in.clone();
        int roundKey[][] = key.clone();
        log.info("open text = ");
        blockToString(in,true);
        log.info("key = ");
        blockToString(key,true);
        state = addRoundKey(state,key);
        log.info("after addRoundKey:");
        blockToString(state,true);

        for (int i=0;i<9;i++){
            log.info("round #"+(i+1)+":");
            state = subBytes(state);
            log.info("after subBytes:");
            blockToString(state,true);

            state = shiftRows(state);
            log.info("after shiftRows:");
            blockToString(state,true);

            state = mixColumns(state);
            log.info("after mixColumns:");
            blockToString(state,true);

            roundKey = keyExpansion(roundKey,i);
            log.info("round #"+(i+1)+" key=");
            blockToString(roundKey,true);

            state = addRoundKey(state, roundKey);
            log.info("after addRoundKey:");
            blockToString(state,true);

        }
        log.info("last transform:");
        state = subBytes(state);
        log.info("after subBytes:");
        blockToString(state,true);

        state = shiftRows(state);
        log.info("after shiftRows:");
        blockToString(state,true);

        roundKey = keyExpansion(roundKey,9);
        state = addRoundKey(state,roundKey);
        log.info("after addRoundKey:");
        blockToString(state,true);
        
        return state;
    }

    /**
     * XOR'им текст с ключем
     * @param state текст
     * @param roundKey ключ
     * @return ___
     */
    protected static int[][] addRoundKey(int[][] state, int[][] roundKey){
        int[][] rezult = state.clone();
        for (int i=0;i<SIZE;i++){
            for(int j=0;j<SIZE;j++){
                rezult[i][j] ^= roundKey[j][i];
            }
        }
        return rezult;
    }

    /**
     * Производит замену элементов в state по таблице invSBox
     * @param state ___
     * @return ___
     */
    private static int[][] subBytes(int [][] state){
        int [][] result = new int[SIZE][SIZE];
        for (int i=0;i<SIZE;i++){
            for (int j=0;j<SIZE;j++){
                int row = getMajorTetrad(state[i][j]);
                int col = getMinorTetrad(state[i][j]);
                result[i][j] = sBox[row][col];
            }
        }
        return result;
    }

    /**
     * Возвращает старшую тетраду младшего байта
     * @param num число, из которого берется младший байт
     * @return ___
     */
    protected static int getMajorTetrad(int num){
        return ((num >> 4) & 0xf);
    }

    /**
     * Возвращает младшую тетраду младшего байта
     * @param num число, из которого берется младший байт
     * @return ___
     */
    protected static int getMinorTetrad(int num){
        return (num & 0xf);
    }

    /**
     * Циклический сдвиг строки массива в зависимости от ее номера влево
     * 0 строка на 0 влево
     * 1 строка на 1 влево
     * и т.д.
     * @param state массив, в котором происходит сдвиг
     * @return измененный массив
     */
    private static int[][] shiftRows(int[][] state){
        int [][] result = new int[SIZE][SIZE];
        for (int i=0;i<SIZE;i++){
            for (int j=0;j<SIZE;j++){
                result[i][j] = state[i][(i+j)%SIZE];
            }
        }
        return result;
    }

    /**
     * умножение по столбцам на фиксированную матрицу
     * @param state массив, который умножаем
     * @return ___
     */
    private static int[][] mixColumns (int[][] state){
        int [][] result = new int[SIZE][SIZE];
        for (int j =0; j <SIZE; j++){
//            result[0][j] = x02(state[0][j]) ^ x03(state[1][j]) ^ state[2][j] ^ state[3][j];
//            result[1][j] = state[0][j] ^ x02(state[1][j]) ^ x03(state[2][j]) ^ state[3][j];
//            result[2][j] = state[0][j] ^ state[1][j] ^ x02(state[2][j]) ^ x03(state[3][j]);
//            result[3][j] = x03(state[0][j]) ^ state[1][j] ^ state[2][j] ^ x02(state[3][j]);

            result[0][j] = gmul(state[0][j],0x02) ^ gmul(state[1][j],0x03) ^ state[2][j] ^ state[3][j];
            result[1][j] = state[0][j] ^ gmul(state[1][j],0x02) ^ gmul(state[2][j],0x03) ^ state[3][j];
            result[2][j] = state[0][j] ^ state[1][j] ^ gmul(state[2][j],0x02) ^ gmul(state[3][j],0x03);
            result[3][j] = gmul(state[0][j],0x03) ^ state[1][j] ^ state[2][j] ^ gmul(state[3][j],0x02);
        }
        return result;
    }


    /**
     * Multiply two numbers in the GF(2^8) finite field defined
     * by the polynomial x^8 + x^4 + x^3 + x + 1
     * Заимствовано у википедии
     * en.wikipedia.org/wiki/Finite_field_arithmetic
     * @param a первый множитель
     * @param b второй множитель
     * @return результат умножения
     */
    protected static int gmul(int a, int b) {
        int p = 0;
        int counter;
        int hi_bit_set;
        for(counter = 0; counter < 8; counter++) {
            if((b & 1) != 0) {
                p ^= a;
            }
            hi_bit_set = (a & 0x80);
            a <<= 1;
            if(hi_bit_set != 0)
                a ^= 0x11b; /* x^8 + x^4 + x^3 + x + 1 */
            b >>= 1;
        }
        return p;
    }

    /**
     * умножение на 0x2 для shiftRows
     * @param num умножаемое число
     * @return ___
     * @deprecated нашел более универсальную замену -
     * {@linkplain Encoder#gmul}
     */
    protected static int x02 (int num){
//        проверка левого бита (если он равен 1b)
//        log.debug("x02 num    = "+Integer.toHexString(num));
        if ((num & 0x80) == 0x80){
            num =  num << 1;
            num =  num ^ 0x1B;
        }else{
            num =  num << 1;
        }
        num =  num & 0xff;
//        log.debug("x02 result = "+Integer.toHexString(num));
        return num;
    }

    /**
     * умножение на 0x3 для shiftRows
     * @param num умножаемое число
     * @return ___
     * @deprecated нашел более универсальную замену -
     * {@linkplain Encoder#gmul}
     */
    protected static int x03(int num){
//        log.debug("x03 num    = "+Integer.toHexString(num));
        num =  x02(num) ^ num;
//        log.debug("x03 result = "+Integer.toHexString(num));
        return num;
    }

    /**
     * Вычисление ключа для нового раунда
     * @param key ключ текущего раунда
     * @param roundNum номер раунда
     * @return новый ключ
     */
    protected static int[][] keyExpansion(int[][] key, int roundNum){
        int[] lastRow = rotWord(key);
//        log.debug("key expansion:");
//        log.debug("after rotWord:");
//        blockToString(key);
        lastRow = subBytesInRow(lastRow);
//        log.debug("after subBytesInRow:");
//        blockToString(key);
        if (roundNum < 8){
            lastRow[0] = lastRow[0]^(1 << roundNum);
//            log.debug("Rcon"+roundNum+"= "+Integer.toHexString(1 << roundNum));
        }else if(roundNum == 8){
            lastRow[0] = lastRow[0]^(0x1b);
//            log.debug("Rcon"+roundNum+"= "+Integer.toHexString(0x1b));
        }else if(roundNum == 9){
            lastRow[0] = lastRow[0]^(0x36);
//            log.debug("Rcon"+roundNum+"= "+Integer.toHexString(0x36));
        }
//        log.debug("after XOR with Rcon:"+Integer.toHexString(lastRow[0]));
        int result[][] = new int[SIZE][SIZE];
//        заполнение происходит построчно
        for(int i=0;i<SIZE;i++){
            result[0][i] = lastRow[i] ^ key[0][i];
            result[1][i] = result[0][i] ^ key[1][i];
            result[2][i] = result[1][i] ^ key[2][i];
            result[3][i] = result[2][i] ^ key[3][i];
        }
        return result;
    }

    /**
     * циклический сдвиг последней строки на один элемент вверх
     * @param key матрица, в которой производится сдвиг
     * @return ___
     */
    private static int[] rotWord(int[][] key){
        int[] result = new int[SIZE];
        for (int i=0;i<SIZE;i++){
            result[i] = key[SIZE-1][(i+1)%SIZE];
        }
        return result;
    }

    /**
     * версия {@linkplain Encoder#subBytes} применимая к столбцу
     * @param col ___
     * @return ___
     */
    private static int[] subBytesInRow(int[] col){
        int [][] switcher = new int[SIZE][SIZE];
        int [] result = new int[SIZE];
        for (int i=0;i<SIZE;i++){
            switcher[i][SIZE-1] = col[i];
        }
        switcher = subBytes(switcher);
        for (int i=0;i<SIZE;i++){
            result[i] = switcher[i][SIZE-1];
        }
        return result;
    }
}
